#!python3.9
# -*- coding: utf-8 -*-
# @Time    : 2025/2/17 10:42
# @Author  : fy
# @File    : myexecutor.py
# @Software: PyCharm
import time
import logging

class MyDriver(object):
    def __init__(self, driver_name, url_dict):
        # TODO to add your driver and implement the query functions
        """
        add your driver
        """
        self.driver_name = driver_name
        logging.debug(self.driver_name)
        url = url_dict.get('url', '')
        if self.driver_name == "postgresql":
            from test.pg_driver import PGDriver as Driver
            url = url_dict['pg_conn']
        elif self.driver_name == "subprocess":
            from test.subcmd_driver import SubcmdDriver as Driver
        else:
            logging.error("driver name not found")
            raise AssertionError("driver not found")
        self.main_driver = Driver(url)
        self.main_driver.initialize(url)


    def initialize(self):
        pass

    def query(self, stmt, timeout=0, return_tuples=False):
        """
        use OrderedDict to return data

        query data like:
            odict_items([('count', 21), ('timetz', datetime.time(14, 49, 32, 899374, tzinfo=datetime.timezone(datetime.timedelta(seconds=28800))))])
        yaml item like:
              metrics:
                    - name: timetz
                      description: current time
                      usage: LABEL
                    - name: info
                      description: info
                      usage: LABEL
                    - name: status
                      description: status desc
                      usage: GAUGE

        query data must same as yaml item
        """
        union_set = set()
        result = self.main_driver.query(stmt, timeout, return_tuples)
        logging.debug("{} query data: {}".format(stmt, result))
        if result is None:
            return []
        for row in result:
            if return_tuples:
                union_set.add(tuple(row))
            else:
                logging.debug(row)
                logging.debug(row.items())
                union_set.add(tuple(row.items()))
        if return_tuples:
            return list(union_set)
        else:
            # Get back to the dict-based format.
            ret = []
            for row in union_set:
                dict_based_row = {}
                for k, v in row:
                    dict_based_row[k] = v
                ret.append(dict_based_row)
            return ret

